package com.smart.thrift.discoveryandbalance.balancer;

import cn.hutool.core.util.ReflectUtil;
import com.google.common.collect.Lists;
import com.smart.thrift.discoveryandbalance.discovery.ServerListUpdater;
import com.smart.thrift.discoveryandbalance.discovery.ThriftServerNode;
import com.smart.thrift.discoveryandbalance.discovery.ThriftServerNodeList;
import com.smart.thrift.discoveryandbalance.discovery.zookeeper.ThriftZookeeperServerListUpdater;
import com.smart.thrift.discoveryandbalance.discovery.zookeeper.ThriftZookeeperServerNode;
import com.smart.thrift.discoveryandbalance.discovery.zookeeper.ThriftZookeeperServerNodeList;
import com.smart.thrift.discoveryandbalance.rule.IRule;
import lombok.extern.slf4j.Slf4j;

import java.util.*;

@Slf4j
public class ThriftZookeeperServerListLoadBalancer extends AbstractServerListLoadBalancer {

    private final ThriftZookeeperServerNodeList serverNodeList;
    private final IRule rule;

    private volatile ServerListUpdater serverListUpdater;

    private final ServerListUpdater.UpdateAction updateAction = this::updateListOfServers;


    public ThriftZookeeperServerListLoadBalancer(ThriftZookeeperServerNodeList serverNodeList, IRule rule) {
        this.serverNodeList = serverNodeList;
        this.rule = rule;
        this.serverListUpdater = new ThriftZookeeperServerListUpdater();
        this.startUpdateAction();
    }

    @Override
    public ThriftZookeeperServerNode chooseServerNode(String key) {
        if (rule == null) {
            return null;
        } else {
            ThriftServerNode serverNode;
            try {
                serverNode = rule.choose(key);
            } catch (Exception e) {
                log.warn("LoadBalancer [{}]:  Error choosing server for key {}", getClass().getSimpleName(), key, e);
                return null;
            }

            if (serverNode instanceof ThriftZookeeperServerNode) {
                return (ThriftZookeeperServerNode) serverNode;
            }
        }

        return null;
    }

    @Override
    public ThriftServerNodeList getThriftServerNodeList() {
        return this.serverNodeList;
    }

    private synchronized void startUpdateAction() {
        log.info("Using serverListUpdater {}", serverListUpdater.getClass().getSimpleName());
        if (serverListUpdater == null) {
            serverListUpdater = new ThriftZookeeperServerListUpdater();
        }
        this.serverListUpdater.start(updateAction);
    }

    public void stopServerListRefreshing() {
        if (serverListUpdater != null) {
            serverListUpdater.stop();
        }
    }

    private void updateListOfServers() {
        Map<String, LinkedHashSet<ThriftZookeeperServerNode>> thriftZookeeperServers = this.serverNodeList.refreshThriftServers();
        Map<ThriftServerNode, Object> poolMap = (Map<ThriftServerNode, Object>) ReflectUtil.getFieldValue(transportKeyedObjectPool, "poolMap");
        Set<ThriftZookeeperServerNode> refreshThriftServers = new LinkedHashSet<>();
        for (LinkedHashSet<ThriftZookeeperServerNode> values : thriftZookeeperServers.values()) {
            refreshThriftServers.addAll(values);
        }
        log.info("refresh nodes : {}, poolKeyList: {}", refreshThriftServers, poolMap.keySet());
        if (Objects.nonNull(poolMap) && !poolMap.isEmpty()) {
            poolMap.keySet().removeAll(refreshThriftServers);
            log.info("some node has bean removed : {}", poolMap.keySet());
            new HashSet<>(poolMap.keySet()).forEach(transportKeyedObjectPool::clear);
        }

        List<String> serverList = Lists.newArrayList();
        StringBuilder sb = new StringBuilder();
        for (Map.Entry<String, LinkedHashSet<ThriftZookeeperServerNode>> serverEntry : thriftZookeeperServers.entrySet()) {
            serverList.add(
                    sb.append(serverEntry.getKey())
                            .append(": ")
                            .append(serverEntry.getValue())
                            .toString()
            );
            sb.setLength(0);
        }

    }
}
